# _*_ coding:utf-8 _*_
# file: main_multithread.py
# project: stock price prediction
# coder: simon zhang
# version:0.06
# email:simonzhang84@163.com
# update:2022/12/03
# new features: matplotlib canvas and thread


import datetime
import math
import os.path
import sys
import matplotlib as mpl
import matplotlib.pyplot as plt
from matplotlib.dates import WeekdayLocator
import numpy as np
import tushare as ts
from PyQt5.QtGui import QIcon
from PyQt5.QtWidgets import QWidget, QProgressBar, QCheckBox, QPushButton, QApplication, QLabel, QLineEdit, QLCDNumber, \
    QComboBox, QFormLayout, QHBoxLayout, QMessageBox
from PyQt5.QtCore import QThread, pyqtSignal
from matplotlib.backends.backend_qt5agg import FigureCanvasQTAgg as Figure
from keras.layers import Dense, LSTM
from keras.models import Sequential
from sklearn.preprocessing import MinMaxScaler
from pandas import DataFrame


def plotResult(widget, data):
        # 显示中文的全局设置
        mpl.rcParams['font.family'] = 'SimHei'
        plt.rcParams['axes.unicode_minus'] = False
        # 画图
        plt.clf()
        widget.figure = plt.figure(f'明日股价预测', facecolor='lightgreen', figsize=(16, 8))
        ax = plt.axes()

        ax.xaxis.set_major_locator(WeekdayLocator())
        ax.xaxis.set_minor_locator(plt.AutoLocator())
        ax.set_title(f'{widget.stock_code}明日股价预测')
        plt.grid(True, 'major', 'both')
        plt.xlabel('时间', fontsize=14)
        plt.ylabel('收盘价', fontsize=14)
        # plt.plot(train['close']
        ax.plot(data[['close', 'predictions']])#
        widget.figure.tight_layout()
        widget.figure.add_subplot(ax)
        ax.legend(['真实值', '预测值'], loc='lower left')
        plt.gcf().autofmt_xdate()
        # vbox = QVBoxLayout(self)
        # vbox.addWidget(self.canvas)
        widget.canvas = Figure(widget.figure)
        widget.setWindowTitle(f'{widget.stock_code}明日股价预测')
        widget.canvas.show()


class PredictThread(QThread):
    price_signal = pyqtSignal(tuple)
    valid_signal = pyqtSignal(DataFrame)

    def __init__(self,period, stock_code, kx, save_data_flag):
        super().__init__()
        self.valid = None
        self.tomorrow_prediction = None
        self.tomorrow_change = None
        self.period = period
        self.stock_code = stock_code
        self.kx = kx
        self.save_data_flag = save_data_flag

    def run(self):
        self.tomorrow_prediction, self.tomorrow_change, self.valid = \
            predict_next_price(period=self.period,
                               stock_code=self.stock_code,
                               kx=self.kx,
                               save_data_flag=self.save_data_flag)
        self.price_signal.emit((self.tomorrow_prediction, self.tomorrow_change))
        self.valid_signal.emit(self.valid)

def getStockData(stock_code='002022', kx='D'):
    """初始化Tushare，获取股票历史数据"""
    token = '828005039aef60a8745faffc3173a21fe03dcb4ec0cbabde2411bacb'
    ts.set_token(token)
    pro = ts.pro_api()
    if stock_code.startswith('0') or stock_code.startswith('3'):
        stock_code += '.SZ'
    elif stock_code.startswith('6'):
        stock_code += '.SH'
    else:
        print('check the code you input')
    last_day = (datetime.date.today() - datetime.timedelta(days=1)).strftime('%Y%m%d')
    df = pro.daily(ts_code=stock_code, start='20170101', end=last_day)
    df.set_index('trade_date', inplace=True)
    df = df.reindex(index=df.index[::-1])
    return df


def predict_next_price(period=100, stock_code='002022', kx='D', save_data_flag=True):
    df = getStockData(stock_code, kx)
    # 保存行情数据为Excel表格
    if save_data_flag:
        if not os.path.exists('./data'):
            os.mkdir('./data')
        df.to_excel(f'./data/{stock_code}-{datetime.date.today()}.xlsx')
    data = df.filter(['close'])
    dataset = data.values
    training_data_len = math.ceil(len(dataset) * 0.95)

    scaler = MinMaxScaler(feature_range=(0, 1))
    scaled_data = scaler.fit_transform(dataset)
    # scaled_data
    train_data = scaled_data[0:training_data_len, :]
    x_train = []
    y_train = []
    for i in range(period, len(train_data)):
        x_train.append(train_data[i - period: i, 0])
        y_train.append(train_data[i, 0])
    x_train, y_train = np.array(x_train), np.array(y_train)
    x_train = np.reshape(x_train, (x_train.shape[0], x_train.shape[1], 1))

    model = Sequential()
    model.add(LSTM(50, return_sequences=True, input_shape=(x_train.shape[1], 1)))
    model.add(LSTM(50, return_sequences=False))
    model.add(Dense(25))
    model.add(Dense(1))
    model.compile(optimizer='adam', loss='mean_squared_error')
    model.fit(x_train, y_train, batch_size=1, epochs=1)
    test_data = scaled_data[training_data_len - period:, :]
    x_test = []
    y_test = dataset[training_data_len:, :]
    for i in range(period, len(test_data)):
        x_test.append(test_data[i - period: i, 0])
    x_test = np.array(x_test)
    x_test = np.reshape(x_test, (x_test.shape[0], x_test.shape[1], 1))
    predictions = model.predict(x_test)
    predictions = scaler.inverse_transform(predictions)
    # rmse = np.sqrt(np.mean(predictions - y_test) ** 2)
    # train = data[:training_data_len]
    valid = DataFrame(data[training_data_len:], copy=True)
    valid.loc[:, 'predictions'] = predictions
    last_certain_days = data[period * (-1):].values
    last_certain_days_scaled = scaler.transform(last_certain_days)
    X_test = []
    X_test.append(last_certain_days_scaled)
    X_test = np.array(X_test)
    X_test = np.reshape(X_test, (X_test.shape[0], X_test.shape[1], 1))
    pred_price = model.predict(X_test)
    pred_price = scaler.inverse_transform(pred_price)[0, 0]
    lastday_prediction = predictions[-1, 0]

    tomorrow_change = (pred_price - lastday_prediction) / lastday_prediction * 100
    return pred_price, tomorrow_change, valid


class StockPrediction(QWidget):

    def __init__(self):
        super().__init__()
        self.stocks = None
        self.canvas = None
        self.valid = None
        self.train_percent = 0.9
        self.initUI()

    def initUI(self):
        btn = QPushButton(u'开始计算', self)
        btn.clicked.connect(self.modelShow)
        btn.resize(btn.sizeHint())
        self.setWindowIcon(QIcon('modeltrain.ico'))
        btn.move(150, 150)
        label = QLabel(u'请输入股票代码    ', self)
        labelcb = QLabel(u'请选择周期类型', self)
        self.period_cb = QComboBox(self)
        self.period_cb.addItems([u'日', u'月', u'周', '120', '60', '30', '5'])

        # label.move(50, 50)

        self.input_stock = QComboBox()
        self.stocks = self.getStockList()
        self.input_stock.addItems(self.stocks)
        self.input_stock.setCurrentIndex(len(self.stocks) - 1)
        self.input_stock.setEditable(True)

        # self.input_stock.move(150, 50)

        label2 = QLabel(u'请输入训练期数    ', self)
        # label2.move(50, 20)
        self.input_period = QLineEdit(self)
        self.input_period.setPlaceholderText('60')
        self.input_period.move(150, 20)
        label3 = QLabel(u'预计明日股价              ', self)
        # label3.move(50, 80)
        self.lcd = QLCDNumber(self)
        # self.lcd.move(150, 80)

        self.lcd.setNumDigits(10)
        # self.lcd.setFixedSize(150, 60)
        self.lcd.setMinimumHeight(60)
        self.lcd.display('00.00')

        self.save_data_checkbox = QCheckBox('保存行情数据', self)
        self.save_data_checkbox.setTristate(False)
        self.save_data_checkbox.setChecked(True)

        # self.show_trend_checkbox = QCheckBox('显示趋势图', self)
        # self.show_trend_checkbox.setChecked(True)

        self.change_label = QLabel('变化趋势：  ', self)

        self.progress = QProgressBar(self)
        self.progress.hide()

        self.btn_show_plot = QPushButton('查看趋势')
        self.btn_show_plot.setEnabled(False)
        self.btn_show_plot.clicked.connect(self.showPlot)

        hbox = QHBoxLayout()
        hbox.addWidget(btn)
        hbox.addWidget(self.btn_show_plot)
        layout = QFormLayout(self)
        layout.addRow(labelcb, self.period_cb)
        layout.addRow(label2, self.input_period)
        layout.addRow(label, self.input_stock)
        layout.addRow(label3, self.lcd)
        layout.addRow(self.save_data_checkbox, self.change_label)
        layout.addRow(QLabel(text='              '), hbox)
        layout.addRow(self.progress)

        self.setLayout(layout)

        self.setGeometry(300, 300, 350, 250)
        self.setWindowTitle('Stock Price Prediction')
        self.show()

    def modelShow(self):


        kx = self.period_cb.currentText()
        if kx == u'日':
            kx = 'D'
        elif kx == u'月':
            kx = 'M'
        elif kx == u'周':
            kx = 'W'

        period = self.input_period.text()
        if not period:
            period = 60
        else:
            period = int(period)
        self.stock_code = self.input_stock.currentText()
        if not self.stock_code.startswith(('3', '6', '0')):
            reply = QMessageBox.information(self, 'Message',
                                         "请输入正确的6位数字股票代码", QMessageBox.StandardButton.Ok, QMessageBox.StandardButton.Ok)
        # 动态进度条
        self.progress.show()
        self.progress.setRange(0, 0)
        self.btn_show_plot.setEnabled(False)
        self.change_label.setText('变化趋势：   ')
        self.saveStockList(self.stock_code)
        self.lcd.display('00.00')

        save_data_flag = self.save_data_checkbox.isChecked()

        self.predict_thread = PredictThread(period=period, stock_code=self.stock_code, kx=kx, save_data_flag=save_data_flag)
        self.predict_thread.price_signal.connect(self.updateDisplay)
        self.predict_thread.valid_signal.connect(self.createPlotCanvas)
        self.predict_thread.finished.connect(self.predictThreadFinished)
        self.predict_thread.start()

    def predictThreadFinished(self):
        self.progress.setRange(0, 100)
        self.progress.setValue(100)
        self.btn_show_plot.setEnabled(True)

    def updateDisplay(self, result):
        tomorrow_prediction, tomorrow_change = result
        self.lcd.display(str(tomorrow_prediction))
        if tomorrow_change > 0:
            self.change_label.setText(f'变化趋势：  预计较上个交易日<span style="color:red"><strong>上涨{tomorrow_change:.2f}%</strong></span>')
        elif tomorrow_prediction == 0:

            self.change_label.setText(f'变化趋势：  预计与上个交易日持平')
        else:
            self.change_label.setText(f'变化趋势：  预计较上个交易日<span style="color:green"><strong>下跌{abs(tomorrow_change):.2f}%</strong></span>')

    def createPlotCanvas(self, valid):
        self.valid = valid


    def showPlot(self):
        plotResult(widget=self, data=self.valid)
        # self.showThread.deleteLater()

    @staticmethod
    def getStockList():
        try:
            with open('data/stocks.txt', 'r', encoding='utf8') as f:
                stocks = f.read().split()
        except FileNotFoundError:
            stocks = []
        return stocks

    def saveStockList(self, stock):
        if stock in self.stocks:
            self.stocks.remove(stock)
        self.stocks.append(stock)
        try:
            with open('data/stocks.txt', 'w', encoding='utf8') as f:
                f.write('\n'.join(self.stocks))
        except Exception:
            pass

def main():
    app = QApplication(sys.argv)
    ex = StockPrediction()
    ex.show()
    sys.exit(app.exec())


if __name__ == '__main__':
    main()
